package io.gitee.zhangbinhub.admin.component

import org.springframework.security.oauth2.server.authorization.OAuth2Authorization
import org.springframework.security.oauth2.server.authorization.OAuth2AuthorizationService
import org.springframework.security.oauth2.server.authorization.OAuth2TokenType
import org.springframework.stereotype.Component
import java.util.concurrent.ConcurrentHashMap

@Component
class AuthTokenService : OAuth2AuthorizationService {
    private val tokenPrefixId = "acp_oauth_token_id_"
    private val tokenPrefixValue = "acp_oauth_token_value_"
    private val tokenPrefixClientId = "acp_oauth_token_client_id_"
    private val tokenIdMap: ConcurrentHashMap<String, OAuth2Authorization> = ConcurrentHashMap()
    private val tokenValueMap: ConcurrentHashMap<String, OAuth2Authorization> = ConcurrentHashMap()
    private val clientTokenMap: ConcurrentHashMap<String, MutableSet<OAuth2Authorization>> = ConcurrentHashMap()

    override fun save(authorization: OAuth2Authorization) {
        tokenIdMap["$tokenPrefixId${authorization.id}"] = authorization
        tokenValueMap["$tokenPrefixValue${authorization.accessToken.token.tokenValue}"] = authorization
        clientTokenMap["$tokenPrefixClientId${authorization.registeredClientId}"].let { clientSet ->
            if (clientSet == null) {
                mutableSetOf(authorization)
            } else {
                clientSet.add(authorization)
                clientSet
            }
        }.apply {
            clientTokenMap["$tokenPrefixClientId${authorization.registeredClientId}"] = this
        }
    }

    override fun remove(authorization: OAuth2Authorization) {
        tokenIdMap.remove("$tokenPrefixId${authorization.id}", authorization)
        tokenValueMap.remove("$tokenPrefixValue${authorization.accessToken.token.tokenValue}", authorization)
        clientTokenMap["$tokenPrefixClientId${authorization.registeredClientId}"]?.let { clientSet ->
            clientSet.removeIf { item ->
                item.id == authorization.id
            }
        }
    }

    fun removeByToken(token: String) {
        findByToken(token)?.apply { remove(this) }
    }

    fun removeTokensByClientIdAndUserName(clientId: String, userName: String) {
        findTokensByClientIdAndUserName(clientId, userName).forEach { remove(it) }
    }

    override fun findById(id: String?): OAuth2Authorization? =
        tokenIdMap["$tokenPrefixId$id"]

    fun findByToken(token: String?): OAuth2Authorization? = findByToken(token, null)

    override fun findByToken(token: String?, tokenType: OAuth2TokenType?): OAuth2Authorization? =
        tokenValueMap["$tokenPrefixValue$token"]

    fun findTokensByClientIdAndUserName(clientId: String, userName: String): Collection<OAuth2Authorization> =
        findTokensByClientId(clientId).filter { it.principalName == userName }

    fun findTokensByClientId(clientId: String): Collection<OAuth2Authorization> =
        clientTokenMap["$tokenPrefixClientId$clientId"]?.apply {
            this.removeIf { item ->
                if (item.accessToken.isExpired) {
                    tokenIdMap.remove(item.id)
                    tokenValueMap.remove(item.accessToken.token.tokenValue)
                    true
                } else {
                    false
                }
            }
        } ?: emptySet()
}